/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * @brief load data instruction ut for ascend910B1
 *
 */
#include <gtest/gtest.h>
#include "kernel_operator.h"
#include "lib/matmul/tiling.h"
#include "impl/matmul/modules/matmul_param.h"
#include "impl/matmul/modules/matmul_policy.h"
#include "impl/matmul/modules/resource/cube_in_buffer/cube_in_buffer.h"
#include "impl/matmul/modules/matmul_private_modules.h"

using namespace std;
using namespace AscendC;
using namespace matmul;

namespace {
template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB,
MATMUL_POLICY_DEFAULT_OF(MatmulPolicy)>
class MatmulImpl: MATMUL_IMPORT_MODULE_PRIVATE(CubeInBufferParamsB) {
    MATMUL_ALLOW_USING_PRIVATE(CubeInBufferParamsB);
public:
    using CubeInBufferParamsB::GetMajorCacheNum;
    using CubeInBufferParamsB::GetTotalCacheNum;
    using CubeInBufferParamsB::GetCurKPos;
    using CubeInBufferParamsB::GetCurMNPos;
    using CubeInBufferParamsB::IsTailBlock;

public:
    using VAR_PARAMS =
        typename MatmulParams<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, GetMatmulVersion(MM_CFG)>::PARAMS;

    MatmulImpl() {
        InitVar();
    }

    VAR_PARAMS& GetVar() {
        return var;
    }

    void InitVar() {
        var.tiling_.SetTiling(&tiling);
        var.tpipe_ = &pipe;
    }

    void SetInitParams(int32_t stepN, int32_t stepKb, int32_t baseN, int32_t baseK) {
        tiling.stepN = stepN;
        tiling.stepKb = stepKb;
        tiling.baseN = baseN;
        tiling.baseK = baseK;
        tiling.iterateOrder = 0;
    }

    void SetRuntimeParams(int32_t baseUseN, int32_t baseUseK) {
        var.baseUseN_ = baseUseN;
        var.baseUseK_ = baseUseK;
    }

private:
    TCubeTiling tiling;
    TPipe pipe;
    VAR_PARAMS var;
};
}

class test_cube_in_buffer_params : public testing::Test {
protected:
    void SetUp() {}
    void TearDown() {}

private:
    using A_TYPE = matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, half, false>;
    using B_TYPE = matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, half, false>;
    using C_TYPE = matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float>;
    using BIAS_TYPE = matmul::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, float>;

    MatmulImpl<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CFG_MDL, void> mm;
};

TEST_F(test_cube_in_buffer_params, all_interface) {
    mm.SetInitParams(2, 2, 32, 32);
    ASSERT_EQ(mm.GetMajorCacheNum(), 2);
    ASSERT_EQ(mm.GetTotalCacheNum(), 4);
    ASSERT_EQ(mm.GetCurKPos(0, 3), 0);
    ASSERT_EQ(mm.GetCurMNPos(0, 3), 3);
    mm.SetRuntimeParams(32, 32);
    ASSERT_FALSE(mm.IsTailBlock());
    mm.SetRuntimeParams(32, 16);
    ASSERT_TRUE(mm.IsTailBlock());
}